import io
import math
import os
import threading
import time
import uuid

from PyQt6.QtCore import QThreadPool

from scripts import video_parser, image_utils, db_utils
from scripts.WorkSingle import Worker
from scripts.beans import Storyboard
from scripts.midjourney import midjourney
from global_config import *

clear_localimg = True


# step 00: 创建项目
def create_project(name, type):
    # 创建目录
    set_cur_project(name)
    workflow_path = f'{workflows_dir}{get_cur_project()}\\'
    image_temps = f'{workflow_path}{workflows_img_dir}'
    if not os.path.exists(image_temps):
        os.makedirs(image_temps)
    frames = f'{workflow_path}{frames_relate_dir}'
    if not os.path.exists(frames):
        os.makedirs(frames)

    # 创建项目配置文件
    project_config = db_utils.load_project_config()
    project_config['type'] = type
    project_config['userid'] = str(uuid.uuid4())
    db_utils.save_project_config(project_config)


# step 01
def parse_movie2storyboard(callback):
    def inner():
        print("runing")
        storyboards_path = get_storyboards_path()

        if os.path.exists(storyboards_path):
            storyboards = db_utils.read_workflow(storyboards_path)
        else:
            # 根据字幕文件，提取出关键帧, 并生成分镜数据
            storyboards = video_parser.movie2Storyboard(get_video_path(), get_srt_path(), get_frames_path())

        # 反推提示词
        for i, bean in enumerate(storyboards):
            if bean.promots is None:
                if bean.local_img is not None:
                    try:
                        bean.promots = image_utils.describe_image(bean.local_img)
                    finally:
                        pass
                else:
                    bean.promots = bean.text

                    # 缓存分镜数据
        db_utils.save_wordflow(storyboards_path, storyboards)

    worker = Worker(inner=inner, callback=lambda a: (callback()))
    thread = QThreadPool.globalInstance()
    thread.start(worker)


# step 01-01
def parse_srt2storyboard(callback, word_limit, resplit):
    def inner():
        print("runing")
        storyboards_path = get_storyboards_path()

        if not resplit and os.path.exists(storyboards_path):
            storyboards = db_utils.read_workflow(storyboards_path)
        else:
            # 根据字幕文件，提取出关键帧, 并生成分镜数据
            storyboards = video_parser.movie2Storyboard(get_video_path(), get_srt_path(), get_frames_path())

        # 根据每个分镜的字数限制，拆分分镜
        mergelist = []
        if word_limit > 0:
            temp_board: Storyboard = None
            temp_count = 0
            for index, board in enumerate(storyboards):
                for srt in board.srtbean:
                    temp_count += len(srt['text'])
                if temp_board is None:
                    temp_board = board
                    mergelist.append(temp_board)
                elif temp_count <= word_limit:
                    temp_board.srtbean = temp_board.srtbean + board.srtbean
                else:  # 超过limit
                    temp_count = 0
                    temp_board = board
                    mergelist.append(temp_board)

            pass

        # 缓存分镜数据
        db_utils.save_wordflow(storyboards_path, mergelist)

    worker = Worker(inner=inner, callback=lambda a:(callback()))
    thread = QThreadPool.globalInstance()
    thread.start(worker)


# 合并关键帧的重复图片
def merge_Duplicate_frames(new_frame, frames: list[str]):
    storyboards_path = get_storyboards_path()
    if os.path.exists(storyboards_path):
        storyboards = db_utils.read_workflow(storyboards_path)
        new_frame_path = None
        for bean in storyboards:
            if bean.img_name == new_frame:
                new_frame_path = bean.local_img
                break

        for bean in storyboards:
            if bean.img_name in frames:
                bean.img_name = new_frame
                bean.local_img = new_frame_path

        # 缓存分镜数据
        db_utils.save_wordflow(storyboards_path, storyboards)


def generate_images():
    storyboards_path = get_storyboards_path()
    if os.path.exists(storyboards_path):
        storyboards = db_utils.read_workflow(storyboards_path)
        for i, bean in enumerate(storyboards):
            if i == 2:
                # MJ 生成图片
                bean.trigger_id = midjourney.send_imagine(bean.promots, bean.ref_img)
                # 缓存分镜数据
                db_utils.save_wordflow(storyboards_path, storyboards)


# 下载图片
def sync_md_imgs():
    storyboards_path = get_storyboards_path()
    image_temp_dir = f'{workflows_dir}{get_cur_project()}\\{workflows_img_dir}\\'
    if not os.path.exists(image_temp_dir):
        os.mkdir(image_temp_dir)

    if os.path.exists(storyboards_path):
        storyboards = db_utils.read_workflow(storyboards_path)
        for bean in storyboards:
            if len(bean.aiImgList) > 0 and len(bean.aiImgNameList) > 0:
                image_utils.save_mj_img(bean.aiImgList, bean.aiImgNameList)


if __name__ == '__main__':
    1

    # 创建线程
    # my_thread = threading.Thread(target=generate_images)
    # my_thread = threading.Thread(target=parse_movie2storyboard)
    # my_thread = threading.Thread(target=merge_Duplicate_frames,
    #                              args=("frame-30.jpg", ['frame-30.jpg', 'frame-60.jpg', 'frame-90.jpg']))
    my_thread = threading.Thread(target=upload_mid_images)
    # my_thread = threading.Thread(target=sync_md_imgs)
    # 启动线程
    my_thread.start()
